Skip to content

Deprecate get_logical_axis_rule_from_config#4568

Open
NuojCheng wants to merge 1 commit into
mainfrom
chengnuojin-train-eval-mesh
Open

Deprecate get_logical_axis_rule_from_config#4568
NuojCheng wants to merge 1 commit into
mainfrom
chengnuojin-train-eval-mesh

Conversation

@NuojCheng

@NuojCheng NuojCheng commented Jul 22, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator

Description

This PR deprecates the get_logical_axis_rules_from_config function in sharding.py. In future, we should use utils.sharding.get_logical_axis_rules to grab logical rule in context.

FIXES: b/536927795

Tests

CI test

Checklist

Before submitting this PR, please make sure (put X in square brackets):

  • I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
  • I have necessary comments in my code, particularly in hard-to-understand areas.
  • I have run end-to-end tests tests and provided workload links above if applicable.
  • I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

@gemini-code-assist

gemini-code-assist Bot commented Jul 22, 2026

Copy link
Copy Markdown

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

@NuojCheng
NuojCheng force-pushed the chengnuojin-train-eval-mesh branch from 61f962b to a316124 Compare July 22, 2026 19:03
@codecov

codecov Bot commented Jul 22, 2026 •
edited
Loading

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

@github-actions

github-actions Bot commented Jul 22, 2026

Copy link
Copy Markdown
Contributor

🤖 Hi @NuojCheng, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

@github-actions

github-actions Bot commented Jul 22, 2026

Copy link
Copy Markdown
Contributor

🤖 I'm sorry @NuojCheng, but I was unable to process your request. Please see the logs for more details.

@NuojCheng
NuojCheng force-pushed the chengnuojin-train-eval-mesh branch from a316124 to 9a2e5c6 Compare July 22, 2026 21:00
@NuojCheng
NuojCheng marked this pull request as ready for review July 22, 2026 21:02
@gemini-code-assist

gemini-code-assist Bot commented Jul 22, 2026

Copy link
Copy Markdown

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

Shuwen-Fang
Shuwen-Fang approved these changes Jul 23, 2026
View reviewed changes
Comment thread src/maxtext/models/deepseek.py
from maxtext.utils.sharding import create_sharding
from maxtext.utils.sharding import maybe_shard_with_logical
from maxtext.utils.sharding import get_logical_axis_rules_from_config
from maxtext.utils.sharding import get_logical_axis_rules

@Shuwen-Fang Shuwen-Fang Jul 23, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

do other models need this update?

@NuojCheng NuojCheng Jul 23, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Only models with explicit sharding onboarded have this change. I updated some other files to make sure get_logical_axis_rules are used.

@NuojCheng
NuojCheng force-pushed the chengnuojin-train-eval-mesh branch from 9a2e5c6 to 86cd649 Compare July 23, 2026 03:22
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@gobbleturk gobbleturk gobbleturk approved these changes
@Shuwen-Fang Shuwen-Fang Shuwen-Fang approved these changes
@RissyRan RissyRan Awaiting requested review from RissyRan RissyRan is a code owner
@michelle-yooh michelle-yooh Awaiting requested review from michelle-yooh michelle-yooh is a code owner
@gagika gagika Awaiting requested review from gagika gagika is a code owner
@richjames0 richjames0 Awaiting requested review from richjames0 richjames0 is a code owner
@shralex shralex Awaiting requested review from shralex shralex is a code owner
@suexu1025 suexu1025 Awaiting requested review from suexu1025 suexu1025 is a code owner
@darisoy darisoy Awaiting requested review from darisoy darisoy is a code owner
@igorts-git igorts-git Awaiting requested review from igorts-git igorts-git is a code owner
@parambole parambole Awaiting requested review from parambole parambole is a code owner
@shuningjin shuningjin Awaiting requested review from shuningjin shuningjin is a code owner
@jiangjy1982 jiangjy1982 Awaiting requested review from jiangjy1982 jiangjy1982 is a code owner
@bvandermoon bvandermoon Awaiting requested review from bvandermoon bvandermoon is a code owner
@xibinliu xibinliu Awaiting requested review from xibinliu xibinliu is a code owner
@huytransformer huytransformer Awaiting requested review from huytransformer huytransformer is a code owner
@aireenmei aireenmei Awaiting requested review from aireenmei aireenmei is a code owner
@abhinavclemson abhinavclemson Awaiting requested review from abhinavclemson abhinavclemson is a code owner
@khatwanimohit khatwanimohit Awaiting requested review from khatwanimohit khatwanimohit is a code owner
@vipannalla vipannalla Awaiting requested review from vipannalla vipannalla is a code owner
@SurbhiJainUSC SurbhiJainUSC Awaiting requested review from SurbhiJainUSC SurbhiJainUSC is a code owner
@hengtaoguo hengtaoguo Awaiting requested review from hengtaoguo hengtaoguo is a code owner
@A9isha A9isha Awaiting requested review from A9isha A9isha is a code owner
@dipannita08 dipannita08 Awaiting requested review from dipannita08 dipannita08 is a code owner

At least 2 approving reviews are required to merge this pull request.

Assignees

No one assigned

Labels

gemini-review pull ready

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@NuojCheng @gobbleturk @Shuwen-Fang